using IntervalArithmetic, IntervalRootFinding, Pipe, Interpolations,
  Optim

"""
    bellman(a1::Float64, a0::Float64, agrid::Vector, vold::Vector;
       thetax::Float16 = 0.36, deltax::Float16 = 0.1, betax::Float16 = 0.98,
       A::Float16 = 0.5)

返回贝尔曼方程的值。

# Arguments
- `a1::Float64`:下一期资本存量。通常调用bellman函数就是通过优化函数来得到最优的`a1`。
- `a0::Float64`:本期的资本存量。
- `agrid::Vector`:资产格点。
- `vold::Vector`:值函数格点。

"""
function bellman(a1::AbstractFloat, a0::AbstractFloat, agrid::Vector, vold::Vector;
   thetax::AbstractFloat = 0.36, deltax::AbstractFloat = 0.1,
   betax::AbstractFloat = 0.98, A::AbstractFloat = 0.5)

  # 通过一阶条件获得先求解劳动
  function chfun(h)
    c0 = a0^thetax*h^(1-thetax) - a1 + (1-deltax)*a0
    A*c0/(1-h)-(1-thetax)*a0^thetax*h^(-thetax)
  end
  h0 = @pipe roots(chfun, 0..1) |> _[1].interval.lo
  # uniroot.all(chfun, interval = c(0,1))
  # print(h0)

  c0 = a0^thetax*h0^(1-thetax) - a1 + (1-deltax)*a0

  if c0 < 0
     return -1e10
  end
  if a1 > agrid[length(agrid)]
     return a1^2*(-1e10)
  end

  il = LinearInterpolation(agrid, vold)
  return -(log(c0) + A*log(1-h0) + betax * il(a1)) # 负值的最小化就是bellman的最大化
end

"""
    para_cmp()

把贝尔曼方程的参数压缩成一个。该函数的参数与`bellman`函数相同，只是少了`a1`。
"""
function para_cmp(a0::AbstractFloat, agrid::Vector, vold::Vector; thetax::AbstractFloat = 0.36,
  deltax::AbstractFloat = 0.1, betax::AbstractFloat = 0.98, A::AbstractFloat = 0.5)
  a1 ->  bellman(a1, a0, agrid, vold, thetax = thetax,
                deltax = deltax, betax = betax, A = A)
end

"""
    decise_range(valuefun::Function, agrid::Vector, m0::Int, x...; v0 = -1e10)

    确定最大值函数对应自变量所处的区间.
"""
function decise_range(valuefun::Function, agrid::Vector, m0::Int, x...; v0 = -1e10)
  # triple ax, bx, cx: [ax,bx] bracket the maximum of Bellman eq.
  ax = bx = agrid[1]
  cx = agrid[length(agrid)]
  m = m0 - 1

  # locate ax <= agrid' <= cx that bracket the maximum
  while m < length(agrid)
    m += 1
    v1 = -valuefun(x...)

    # 如果V1一直比V0大，意味着函数的拐点没有达到，所以m还要往前递进
    # 直到V1比V0小，意味着拐点就在m-1和m之间了，此时就可以给cx赋值
    if v1 > v0 # new value at lower bound agrid[0]
      if m==1
        ax = agrid[m]
        bx = agrid[m]
        cx = agrid[m+1]
      elseif m == length(agrid)
        ax = agrid[m-1]
        bx = agrid[m]
        cx = agrid[m]
      else
        ax = agrid[m-1]
        bx = agrid[m]
        cx = agrid[m+1]
      end
      v0 = v1
      m0 = m   # monotonocity of the value function
    else
      m = length(agrid)
    end
  end
  return (ax = ax, bx = bx, cx = cx, m0 = m0, v0 = v0)
end

"""
    searmax(bellman::Function, pos::Tuple; agrid = zeros(10),
      vold = ones(10), eps = 0.01)

包含边界的最大值搜索.
"""
function searmax(bellman::Function, pos::Dict; agrid::Vector = zeros(10),
  vold::Vector = ones(10), eps::Float64 = 0.01)
  il = LinearInterpolation(agrid, vold)
  if pos[:ax] == pos[:bx]
    pos[:bx] = pos[:ax] + eps * (agrid[2] - agrid[1])
    if il(pos[:bx]) < il(pos[:ax])
      aopt = agrid[1]
    else
      aopt = optimize(bellman, pos[:ax],pos[:cx], GoldenSection()).minimizer
    end
  elseif pos[:bx] == pos[:cx]
    pos[:bx] = pos[:cx] - eps * (agrid[end] - agrid[length(agrid)-1])
    if il(pos[:bx]) < il(pos[:cx])
      aopt = agrid[end]
    else
      aopt = optimize(bellman, pos[:ax],pos[:cx], GoldenSection()).minimizer
    end
  else
    aopt = optimize(bellman, pos[:ax],pos[:cx], GoldenSection()).minimizer
  end
  return aopt
end
